// 版权声明
#include "nms_cuda_kernel.cuh"
#include "pytorch_cuda_helper.hpp"

/**
 * @brief CUDA核启动器，用于非最大抑制（NMS）。
 * 
 * @param boxes 候选框的坐标，格式为[y1, x1, y2, x2]。
 * @param scores 候选框的分数。
 * @param iou_threshold NMS算法中使用的IoU（交并比）阈值。
 * @param offset 用于索引调整的偏移量。
 * @return Tensor 包含经过NMS后保留的候选框索引的张量。
 */
Tensor NMSCUDAKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold, int offset) {
  // 确保操作在正确的CUDA设备上进行
  at::cuda::CUDAGuard device_guard(boxes.device());

  // 如果没有候选框，则直接返回一个空张量
  if (boxes.numel() == 0) {
    return at::empty({0}, boxes.options().dtype(at::kLong));
  }

  // 根据分数对候选框进行排序，获取排序后的索引
  // std::get 是 C++ 标准库中的一个函数，它主要用于访问 std::tuple 或者其他聚合类型（如 std::pair）中的元素。
  auto order_t = std::get<1>(scores.sort(0, /*descending=*/true));
  // 使用排序后的索引对候选框进行排序
  auto boxes_sorted = boxes.index_select(0, order_t);

  // 计算候选框的数量
  int boxes_num = boxes.size(0);
  // 计算列块数，用于CUDA网格和块的配置
  const int col_blocks = (boxes_num + threadsPerBlock - 1) / threadsPerBlock;
  const int col_blocks_alloc = GET_BLOCKS(boxes_num, threadsPerBlock);
  // 创建一个掩码张量，用于存储NMS过程中每个候选框的状态
  Tensor mask = at::empty({boxes_num, col_blocks}, boxes.options().dtype(at::kLong));
  // 配置CUDA网格和块的尺寸
  dim3 blocks(col_blocks_alloc, col_blocks_alloc);
  dim3 threads(threadsPerBlock);
  // 获取当前的CUDA流
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  // 启动NMS CUDA核
  nms_cuda<<<blocks, threads, 0, stream>>>(boxes_num, iou_threshold, offset, boxes_sorted.data_ptr<float>(), (unsigned long long*)mask.data_ptr<int64_t>());
  
  // 创建一个张量，用于存储最终保留的候选框索引
  at::Tensor keep_t = at::zeros({boxes_num}, boxes.options().dtype(at::kBool).device(at::kCUDA));
  // 从掩码中收集保留的候选框索引
  gather_keep_from_mask<<<1, min(col_blocks, THREADS_PER_BLOCK), col_blocks * sizeof(unsigned long long), stream>>>(keep_t.data_ptr<bool>(), (unsigned long long*)mask.data_ptr<int64_t>(), boxes_num);
  // 检查CUDA操作是否有错误
  AT_CUDA_CHECK(cudaGetLastError());
  // 返回排序后的索引中对应保留候选框的索引
  return order_t.masked_select(keep_t);
}